import torch
import cv2
import numpy as np
# 加载模型
#model = torch.hub.load('/home/agilex/catkin_ws/src/scout_deeplearning/yolov5_autopilot','custom', '../yolov5/yolov5s.pt', source='local')
print(torch.cuda.is_available())
path="/home/agilex/catkin_ws/src/scout_deeplearning/yolov5_pedestrian"
model = torch.hub.load(path,'custom', path+'/yolov5s.pt', source='local',force_reload=True)
#model=torch.load("best.pt")
# 设置模型为推理模式
model.eval()
# result=model("001.jpg")
# result.show()
#连接摄像头
cap = cv2.VideoCapture(0)  # 0代表默认摄像头，如果有多个摄像头请根据实际情况修改
ret, frame = cap.read()
print(frame.shape)
#调色板
pose_palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255],
                                      [153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255],
                                      [255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],
                                      [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]],
                                     dtype=np.uint8)
palette_count=len(pose_palette)
while True:
    # 读取摄像头的每一帧
    ret, frame = cap.read()
    # 将图像转换为RGB格式
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # 使用Yolov5进行目标检测
    results = model(frame_rgb)
    # 获取检测结果的边界框、类别和置信度
    boxes = results.xyxy[0].numpy()[:, :4]  # 边界框
    boxes=boxes.astype(int)# 将小数转换为整数
    scores = results.xyxy[0].numpy()[:, 4]  # 置信度
    classes = results.xyxy[0].numpy()[:, 5]  # 类别
    # 遍历每个检测结果
    for (x, y, w, h), score, cls in zip(boxes, scores, classes):
        # 绘制边界框和类别标签
        #print(x,y,w,h)
        color=tuple(pose_palette[int(cls)%palette_count].tolist())
        cv2.rectangle(frame, (x, y), (w, h), color, 2)
        text=f'{model.names[int(cls)]} {score:.2f}'
        (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)#获取字体的宽度和高度
        box_coords = ((x, y - text_height), (x + text_width, y))
        cv2.rectangle(frame, box_coords[0], box_coords[1], color, cv2.FILLED)#给文字添加背景
        cv2.putText(frame, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
                    (255, 255, 255), 2)

    # 显示检测结果图像
    cv2.imshow('Yolov5s', frame)
    # 按下 'q' 键退出程序
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
# 释放摄像头资源并关闭窗口
cap.release()
cv2.destroyAllWindows()